import torch
from random import random
from trainNut import NeuralNetwork, device, test_data

model = NeuralNetwork().to(device)
model.load_state_dict(torch.load("cnnModel.pth"))

classes = [
    "T-shirt/top",
    "Trouser",
    "Pullover",
    "Dress",
    "Coat",
    "Sandal",
    "Shirt",
    "Sneaker",
    "Bag",
    "Ankle boot",
]
print(f'labels: {classes}')

randomI = int(random() * len(test_data))
if randomI >= len(test_data)-10:
    randomI = len(test_data)-11

model.eval()
for i in range(10):
    x, y = test_data[i+randomI][0], test_data[i+randomI][1]
    x=x.unsqueeze(0) # x的形状[1,28,28] 模型第一层输入是conv2d，吃4D输入，利用unsqueeze加多一个维度变成[1,1,28,28]
    with torch.no_grad():
        x = x.to(device)
        pred = model(x)
        predicted, actual = classes[pred[0].argmax(0)], classes[y]
        print(f'Predicted: {predicted:>11s}, Actual: {actual:>11s}, {predicted==actual}',end='')
        if predicted!=actual:
            print(pred)
        else:
            print('\n',end='')
        